import pylab as plt
from params import *
from fit import *
from plfit import *

class Plotter:
    def __init__(self,values):
        self.data = values
        print 'Plotter with %d values'%(len(self.data))
        self._stats()

    def _stats(self):
        v = {}
        for x in self.data:
            v.setdefault(x,0)
            v[x] += 1

        self.x = sorted(v)
        self.y = [v[k] for k in self.x]
        c = [0.0]
        for a in self.y:
            c.append(c[-1]+a)
        tot = c[-1]
        self.c = [1.0-i/tot for i in c[:-1]]
        self.y = [1.0*i/tot for i in self.y]

    def fit_plot(self,filename,log=True,**kwargs):
        if log:
            #p = plfit(self.data)
            #alpha = p._alpha
            #xmin = p._xmin
            p1 = PowerLawFit(self.data)
            alpha, xmin = p1.fit(**kwargs)
            pw_c_x,pw_c_y = p1.fitted_ccdf()
            pw_p_x,pw_p_y = p1.fitted_pdf()
            try:
                imin = self.x.index(xmin)
            except:
                imin = len(self.x)-1
            real_pdf = self.y[imin]
            real_cdf = self.c[imin]

            c_norm = real_cdf/pw_c_y[0]
            pw_c_y *= c_norm

            p_norm = real_pdf/pw_p_y[0]
            pw_p_y *= p_norm


        p2 = LogNormalFit(self.data)
        mu,sigma = p2.fit()
        ln_p_x,ln_p_y = p2.fitted_pdf()
        ln_c_x,ln_c_y = p2.fitted_ccdf()

        legs = ['Data']
        plt.figure()
        plt.clf()
        plt.axes(FIG_AXES2)
        plt.loglog(self.x,self.c,'cx')
        if log:
            plt.loglog(pw_c_x, pw_c_y, 'k-',linewidth=1)
            legs.append("Power-law")
        plt.loglog(ln_c_x, ln_c_y, 'k--',linewidth=1)
        legs.append("Log-normal")
        xmin,xmax = min(self.x),max(self.x)
        ymin,ymax = min(self.c),max(self.c)
        plt.axis([xmin,xmax,ymin,ymax])
        plt.grid(True)
        plt.ylabel('CCDF')
        plt.legend(legs, loc='lower left',numpoints=1)
        plt.savefig('%s_ccdf.pdf'%(filename))
        plt.close()

        plt.figure()
        plt.clf()
        plt.axes(FIG_AXES2)
        plt.loglog(self.x,self.y,'cx')
        if log:
            plt.loglog(pw_p_x, pw_p_y, 'k-',linewidth=1)
            legs.append("Power-law")
        plt.loglog(ln_p_x, ln_p_y, 'k--',linewidth=1)
        legs.append("Log-normal")
        xmin,xmax = min(self.x),max(self.x)
        ymin,ymax = min(self.y),max(self.y)
        plt.axis([xmin,xmax,ymin,ymax])
        plt.grid(True)
        plt.ylabel('PDF')
        plt.legend(legs,loc='lower left',numpoints=1)

        plt.savefig('%s_pdf.pdf'%(filename))
        plt.close()


if __name__ == '__main__':
    a = 2.0
    import pylab as plt
    import random

    #R = map(lambda i: int(i*1000),powerlaw.rvs(a, size=1000))
    values  = [int(random.paretovariate(a)) for i in range(10000)]
    print min(values), max(values)


    p = Plotter(values)
    p.fit_plot('prova.pdf',log=True)
